/*
* Description: 驱动打点
*/
#include "driver_report.h"
#include <cstdio>
#include <cstring>
#include <set>
#include "log.h"
#include "text_util/string_proc.h"
#include "../common/guest_utils.h"
#include "../common/text_utils.h"
#include "../common/session_utils.h"

using namespace COMMON;
namespace {
    const int HOLD_TIME = 60 * 1000;
}
namespace DFX {
void DriverReport::Init(std::shared_ptr<WINCONTROL::SendToHost> sendtoHost)
{
    if (sendtoHost == nullptr) {
        LOG_ERR("[DFX - DriverReport] - sendtoHost is null!");
    }
    m_sendtoHost = sendtoHost;
}

void DriverReport::Deinit()
{
    if (m_sendtoHost != nullptr) {
        m_sendtoHost.reset();
    }
}

void DriverReport::ReportData()
{
    // 当前在线的驱动
    std::vector<DriverReportInfo> onlineDrivers = GetDriversInfo();
    if (onlineDrivers.empty()) {
        //查询为空，过一分钟再查一次，避免打点异常上报
        LOG_INFO("[DFX - DriverReport] GetDriversInfo empty");
        Sleep(HOLD_TIME);
        LOG_INFO("[DFX - DriverReport] GetDriversInfo again");
        onlineDrivers = GetDriversInfo();
    }
    // 驱动丢失上报
    FillReportLostData(onlineDrivers);
    // 驱动状态异常上报
    FillReportStatusData(onlineDrivers);
    LOG_INFO("[DFX - DriverReport] report");
}

void DriverReport::CheckInstallLostDrivers()
{
    if (!SessionUtils::IsUserLogon()) {
        return;
    }
    std::vector<DriverReportInfo> onlineDrivers = GetDriversInfo();
    std::vector<std::string> lostDrivers2 = GetLostDriversInfo(onlineDrivers);
    if (lostDrivers2.empty()) {
        LOG_INFO("InstallLostDrivers success");
    } else {
        std::ostringstream oss;
        for (size_t i = 0; i < lostDrivers2.size(); ++i) {
            oss << lostDrivers2[i];
            if (i != lostDrivers2.size() - 1) oss << ", ";
        }
        LOG_INFO("InstallLostDrivers failed: {}", oss.str());
        InstallLostDrivers(lostDrivers2);
    }
}

void DriverReport::FillReportLostData(const std::vector<DriverReportInfo> onlineDrivers)
{
    // 丢失的驱动
    std::vector<std::string> lostDrivers = GetLostDriversInfo(onlineDrivers);
    // 打点上报
    RecordLostBiLog(lostDrivers);
    // 安装丢失的驱动
    if (!lostDrivers.empty()) {
        LOG_INFO("[DFX - DriverReport]LostDrivers: {}", fmt::join(lostDrivers, ", "));
        InstallLostDrivers(lostDrivers);
        CheckInstallLostDrivers();
        LOG_INFO("show update_confirm restore to choose");
        std::string updateCmd = "C:\\ProgramData\\VMAgent\\update\\update_confirm.exe restore";
        int select = GuestUtils::ExecCmdResult(updateCmd, true);
        if (select == 1) {
            LOG_INFO("choose Yes.restore then reboot");
            GuestUtils::ExecCmd("shutdown -r -t 0", false);
        } else if (select == 0) {
            LOG_INFO("user cancel restore");
        } else {
            LOG_INFO("user no operation or something,stop update_confirm");
            GuestUtils::ExecCmd("cmd /q /C \"taskkill /F /IM update_confirm.exe\"", false);
        }
    }
}

void DriverReport::InstallLostDrivers(const std::vector<std::string> &lostDrivers)
{
    for (std::string element : lostDrivers) {
        auto cmd = cmdMap.find(element);
        if (cmd != cmdMap.end()) {
            std::string batPath = element + ".bat";
            if (TextUtils::Cmds2File(cmd->second, batPath)) {
                GuestUtils::ExecCmd(batPath + " 1>>C:\\ProgramData\\VMAgent\\log\\update.log 2>&1", false);
            } else {
                LOG_ERR("failed to generate {}", batPath);
                continue;
            }
        } else {
            auto it = infMap.find(element);
            if (it != infMap.end()) {
                GuestUtils::ExecCmd("cmd /q /C \"echo %date% %time% - install lost driver"
                    " >>C:\\ProgramData\\VMAgent\\log\\update.log 2>&1"
                    " && PnpUtil /add-driver " + BASE_PATH + it->second + " /install"
                    " >>C:\\ProgramData\\VMAgent\\log\\update.log 2>&1\"", false);
            }
        }
    }
    GuestUtils::ExecCmd("pnputil.exe /scan-devices /async", false, false);
}

void DriverReport::RecordLostBiLog(const std::vector<std::string> &lostDrivers)
{
    // 上报
    if (!m_sendtoHost) {
        LOG_ERR("[DFX - DriverReport] sendtoHost is null!");
        return;
    }
    for (std::string element : lostDrivers) {
        DriverReportInfo driverInfo;
        driverInfo.driverName = element;
    }
}

std::string DriverReport::ToJsonString(const DriverReportInfo &driverReportInfo)
{
    if (driverReportInfo.driverName.empty()) {
        return "";
    }
    std::ostringstream reportStream;
    reportStream << "{\"driverName\":\"" << driverReportInfo.driverName;
    if (!driverReportInfo.driverVersion.empty()) {
        reportStream << "\",\"driverVersion\":\"" << driverReportInfo.driverVersion;
    }
    if (!driverReportInfo.driverErrCode.empty()) {
        reportStream << "\",\"driverErrCode\":\"" << driverReportInfo.driverErrCode;
    }
    reportStream << "\"}";
    return reportStream.str();
}

std::vector<DriverReportInfo> DriverReport::GetDriversInfo()
{
    std::vector<DriverReportInfo> driverReportInfos;
    std::array<char, ARRAY_SIZE> buffer;
    std::string result;
    std::unique_ptr<FILE, decltype(&_pclose)> pipe(
        _popen("PnPUtil /enum-drivers", "r"), _pclose);

    if (!pipe) {
        LOG_ERR("[DFX - DriverReport] popen() failed!");
    }

    while (fgets(buffer.data(), buffer.size(), pipe.get()) != nullptr) {
        result += buffer.data();
    }

    // 解析输出
    std::istringstream iss(result);
    std::string line;
    std::string label;
    DriverReportInfo info;

    while (std::getline(iss, line)) {
        if (line.find(PUBLISH_NAME_KEY) != std::string::npos) {
            std::istringstream stream(line);
            stream >> label >> info.publishName;
        }
        if (line.find(ORIGIN_NAME_KEY) != std::string::npos) {
            std::istringstream stream(line);
            stream >> label >> info.driverName;
        }
        if (line.find(DRIVER_VERSION_KEY) != std::string::npos) {
            std::istringstream stream(line);
            stream >> label >> label >> info.driverVersion;
        }
        if (!info.publishName.empty() && !info.driverName.empty() && !info.driverVersion.empty()) {
            driverReportInfos.push_back(info);
            info = {};
        }
    }

    return driverReportInfos;
}

std::vector<std::string> DriverReport::GetLostDriversInfo(const std::vector<DriverReportInfo> &onlineDrivers)
{
    std::vector<std::string> allDriverNames = DRIVER_NAMES;
    for (auto it = allDriverNames.begin(); it != allDriverNames.end();) {
        bool find = false;
        for (auto info : onlineDrivers) {
            if (info.driverName == *it) {
                find = true;
                break;
            }
        }
        if (find) {
            it = allDriverNames.erase(it);
        } else {
            ++it;
        }
    }
    return allDriverNames;
}

void DriverReport::FillReportStatusData(const std::vector<DriverReportInfo> onlineDrivers)
{
    // 驱动状态
    std::vector<DriverReportInfo> driversStatusInfo = GetDriversStatusInfo();
    // 驱动状态异常列表
    std::vector<DriverReportInfo> driversStatusErrInfo = GetDriversErrorInfo(onlineDrivers, driversStatusInfo);
    // 打点上报
    RecordStatusBiLog(driversStatusErrInfo);
}

void DriverReport::RecordStatusBiLog(const std::vector<DriverReportInfo> &driversStatusInfo)
{
    // 上报
    if (!m_sendtoHost) {
        LOG_ERR("[DFX - DriverReport] sendtoHost is null!");
        return;
    }
    for (auto element : driversStatusInfo) {
        std::string data = ToJsonString(element);
        if (data.empty()) {
            continue;
        }
    }
}

std::vector<DriverReportInfo> DriverReport::GetDriversStatusInfo()
{
    std::vector<DriverReportInfo> driverReportInfos;
    std::array<char, ARRAY_SIZE> buffer;
    std::string result;
    std::unique_ptr<FILE, decltype(&_pclose)> pipe(
        _popen("PnPUtil /enum-devices", "r"), _pclose);

    if (!pipe) {
        LOG_ERR("[DFX - DriverReport] popen() failed!");
    }

    while (fgets(buffer.data(), buffer.size(), pipe.get()) != nullptr) {
        result += buffer.data();
    }

    // 解析输出
    std::istringstream iss(result);
    std::string line;
    std::string label;
    DriverReportInfo info;

    while (std::getline(iss, line)) {
        if (line.find(DRIVER_NAME_KEY) != std::string::npos) {
            std::istringstream stream(line);
            stream >> label >> info.driverName;
        }
        if (line.find(DRIVER_STATUS_KEY) != std::string::npos) {
            std::istringstream stream(line);
            stream >> label >> info.driverStatus;
        }
        if (line.find(DRIVER_ERRCODE_KEY) != std::string::npos) {
            std::istringstream stream(line);
            stream >> label >> info.driverErrCode;
        }
        if (!info.driverName.empty() && line.empty()) {
            driverReportInfos.push_back(info);
            info = {};
        }
    }
    if (!info.driverName.empty()) {
        driverReportInfos.push_back(info);
    }

    return driverReportInfos;
}

std::vector<DriverReportInfo> DriverReport::GetDriversErrorInfo(const std::vector<DriverReportInfo> &onlineDrivers,
    const std::vector<DriverReportInfo> &driversStatusInfo)
{
    std::map<std::string, DriverReportInfo> onlineDriverMap;
    std::set<std::string> onlineNameSet;
    for (auto item : onlineDrivers) {
        onlineDriverMap.insert(std::make_pair(item.publishName, item));
        if (item.driverName != VIOGPUDO) {
            onlineNameSet.insert(item.publishName);
        }
    }
    std::vector<std::string> allDriverNames = DRIVER_NAMES;
    std::map<std::string, std::string> allDriverMap;
    for (auto item : allDriverNames) {
        allDriverMap.insert(std::make_pair(item, item));
    }
    std::set<std::string> relatedDriverSet; // related inf
    std::vector<DriverReportInfo> driverErrorInfo;
    for (auto item : driversStatusInfo) {
        if (onlineNameSet.find(item.driverName) != onlineNameSet.end()) {
            onlineNameSet.erase(item.driverName);
            relatedDriverSet.insert(onlineDriverMap[item.driverName].driverName);
        }
        if (item.driverStatus != DRIVER_STATUS_ERR && item.driverStatus != DRIVER_STATUS_DISABLE) {
            continue;
        }
        if (onlineDriverMap.count(item.driverName) == 0) {
            continue;
        }
        DriverReportInfo info = onlineDriverMap[item.driverName];
        if (allDriverMap.count(info.driverName) == 0) {
            continue;
        }
        std::string errCode;
        if (item.driverStatus == DRIVER_STATUS_DISABLE) {
            errCode = DRIVER_ERRCODE_DISABLE;
        } else {
            errCode = item.driverErrCode;
        }
        driverErrorInfo.push_back(AssembleDriverErrorInfo(info, errCode));
    }
    for (auto item : onlineNameSet) {
        DriverReportInfo info = onlineDriverMap[item];
        if (allDriverMap.count(info.driverName) == 0 ||
            relatedDriverSet.find(info.driverName) != relatedDriverSet.end()) {
            continue;
        }
        driverErrorInfo.push_back(AssembleDriverErrorInfo(info, DRIVER_ERRCODE_UNINSTALL));
    }
    return driverErrorInfo;
}

DriverReportInfo DriverReport::AssembleDriverErrorInfo(const DriverReportInfo &info, const std::string &errCode)
{
    DriverReportInfo errInfo;
    errInfo.driverName = info.driverName;
    errInfo.driverVersion = info.driverVersion;
    errInfo.driverErrCode = errCode;
    LOG_ERR("[DFX - DriverReport] DriverStatusError-DriverName: {}, DriverVersion: {}, errCode: {}",
        info.driverName, info.driverVersion, errCode);
    if (strcmp(info.driverName.c_str(), SCREAM.c_str()) == 0
        && strcmp(errCode.c_str(), DRIVER_ERRCODE_UNINSTALL.c_str()) == 0) {
        auto cmd = cmdMap.find(SCREAM);
        std::string batPath = "scream.bat";
        if (TextUtils::Cmds2File(cmd->second, batPath)) {
            GuestUtils::ExecCmd(batPath + " 1>>C:\\ProgramData\\VMAgent\\log\\update.log 2>&1", false);
        } else {
            LOG_ERR("failed to generate batPath");
        }
    }
    return errInfo;
}
}